import os
import xgboost as xgb
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import joblib
from utils.data_loader import load_data
from utils.log import setup_logger
from utils.preprocessor import preprocess_data, save_data
from utils.evaluator import evaluate_model
from sklearn.metrics import classification_report


def main(data_path, seed):
    # 配置路径
    log_path, log_filename = os.path.join('../log', 'training'), 'training'
    fig_path = '../fig'
    model_path = '../model'

    # 设置日志
    logger = setup_logger(log_path, log_filename)
    logger.info("开始建立模型...")
    print("开始建立模型...")

    # 加载数据
    logger.info("加载数据中...")
    print("加载数据中...")
    data = load_data(data_path, 'train.csv')

    # 数据预处理
    logger.info("正在处理数据...")
    print("正在处理数据...")
    data, feature_names = preprocess_data(data)
    # 保存处理后的数据
    save_data(data, '../data', 'train_processed.csv')

    # 划分数据集
    X = data.drop('Attrition', axis=1)
    y = data['Attrition']
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=seed, stratify=y
    )

    # 创建并训练随机森林模型
    rfc_model = RandomForestClassifier(n_estimators=300, max_depth=5, random_state=seed)
    rfc_model.fit(X_train, y_train)

    # 创建并训练 XGBoost模型
    logger.info("正在训练XGBoost模型...")
    print("正在训练XGBoost模型...")
    model = xgb.XGBClassifier(
        # 设置XGBoost分类器的参数
        objective='binary:logistic',  # 目标函数为二分类的对数损失函数
        eval_metric='auc',  # 评估指标为 ROC曲线下的面积
        n_estimators=300,  # 设置估计器（决策树）的数量
        learning_rate=0.1,  # 学习率为 0.05，控制每棵树对最终结果的贡献程度
        max_depth=5,  # 决策树的最大深度为 5，用于控制模型的复杂度
        subsample=0.8,  # 对样本进行随机采样的比例，80%的样本用于训练每棵树，以减少过拟合
        colsample_bytree=0.8,  # 对特征进行随机采样的比例，每棵树使用80%的特征，以增加模型的泛化能力
        random_state=seed,  # 随机种子，确保模型的可重复性
        scale_pos_weight=len(y_train[y_train == 0]) / len(y_train[y_train == 1])  # 正负样本的权重比，处理类别不平衡问题
    )

    model.fit(X_train, y_train)
    # 评估模型
    logger.info("正在评估模型...")
    print("正在评估模型...")
    report = classification_report(y_test, model.predict(X_test))
    roc_auc = evaluate_model(model, X_test, y_test, fig_path, 'xgb', 'train')
    print(f"评估报表: {report}")
    logger.info(f"模型ROC-AUC得分: {roc_auc:.4f}")
    print(f"模型ROC-AUC得分: {roc_auc:.4f}")

    # 保存模型
    model_file = os.path.join(model_path, 'xgb_model.pkl')
    joblib.dump(model, model_file)
    logger.info(f"模型已保存到 {model_file}")
    print(f"模型已保存到 model/xgb_model.pkl")


if __name__ == "__main__":
    seed = 6
    input_data_path = '../data'
    main(input_data_path, seed)
